import torch
import json
import numpy as np
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from torch.nn import Sigmoid

class BERTMultiLabelClassifier:
    """
    BERT多标签文本分类器
    """
    
    def __init__(self, model_path, threshold=0.5):
        """
        初始化分类器
        
        Args:
            model_path (str): 模型路径
            threshold (float): 分类阈值
        """
        self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
        self.threshold = threshold
        
        # 加载模型和分词器
        print(f"正在从 {model_path} 加载模型...")
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_path,
            torch_dtype=torch.float32,
            device_map=None
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        
        # 移动模型到设备
        self.model = self.model.to(self.device)
        self.model.eval()
        
        # 加载标签列表
        label_list_path = f"{model_path}/label_list.json"
        with open(label_list_path, 'r', encoding='utf-8') as f:
            self.label_list = json.load(f)
        
        # 加载标签描述
        categories_path = "./data/categories.json"
        with open(categories_path, 'r', encoding='utf-8') as f:
            self.categories = json.load(f)
        
        print(f"模型加载完成，设备: {self.device}")
        print(f"支持的标签: {self.label_list}")
    
    def predict(self, text, return_probabilities=False):
        """
        对单个文本进行预测
        
        Args:
            text (str): 输入文本
            return_probabilities (bool): 是否返回概率值
        
        Returns:
            dict: 预测结果
        """
        # 文本预处理
        inputs = self.tokenizer(
            text,
            truncation=True,
            padding='max_length',
            max_length=512,
            return_tensors='pt'
        )
        
        # 移动输入到设备
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 模型推理
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
        
        # 计算概率
        sigmoid = Sigmoid()
        probabilities = sigmoid(logits).cpu().numpy()[0]
        
        # 二值化预测
        predictions = (probabilities > self.threshold).astype(int)
        
        # 构建结果
        result = {
            'text': text,
            'predicted_labels': [],
            'label_probabilities': {}
        }
        
        for i, (label, prob, pred) in enumerate(zip(self.label_list, probabilities, predictions)):
            result['label_probabilities'][label] = float(prob)
            if pred == 1:
                result['predicted_labels'].append(label)
        
        if return_probabilities:
            return result
        else:
            return {
                'text': text,
                'predicted_labels': result['predicted_labels']
            }
    
    def predict_batch(self, texts, return_probabilities=False):
        """
        对批量文本进行预测
        
        Args:
            texts (list): 文本列表
            return_probabilities (bool): 是否返回概率值
        
        Returns:
            list: 预测结果列表
        """
        results = []
        for text in texts:
            result = self.predict(text, return_probabilities)
            results.append(result)
        return results
    
    def explain_prediction(self, text):
        """
        解释预测结果
        
        Args:
            text (str): 输入文本
        
        Returns:
            dict: 详细的预测解释
        """
        result = self.predict(text, return_probabilities=True)
        
        explanation = {
            'text': text,
            'predicted_labels': result['predicted_labels'],
            'label_explanations': {}
        }
        
        # 按概率排序
        sorted_labels = sorted(
            result['label_probabilities'].items(),
            key=lambda x: x[1],
            reverse=True
        )
        
        for label, prob in sorted_labels:
            explanation['label_explanations'][label] = {
                'probability': prob,
                'predicted': label in result['predicted_labels'],
                'description': self.categories.get(label, "未知标签")
            }
        
        return explanation

def main():
    """
    演示推理功能
    """
    # 模型路径
    model_path = "./models/bert-multilabel"
    
    try:
        # 初始化分类器
        classifier = BERTMultiLabelClassifier(model_path)
        
        # 测试样例
        test_texts = [
            "绩效奖金缩水了",
            "这阵子加班不多，能接受；工具很好用，从没掉链子",
            "权限申请复杂，卡在审批",
            "职业路径不清晰，晋升机会很少",
            "主管很支持我的工作，经常给予指导",
            "跨部门沟通效率低，信息不透明"
        ]
        
        print("\n=== 批量预测示例 ===")
        results = classifier.predict_batch(test_texts)
        for result in results:
            print(f"文本: {result['text']}")
            print(f"预测标签: {result['predicted_labels']}")
            print("-" * 50)
        
        print("\n=== 详细预测解释示例 ===")
        for text in test_texts[:3]:  # 只展示前3个
            explanation = classifier.explain_prediction(text)
            print(f"\n文本: {explanation['text']}")
            print(f"预测标签: {explanation['predicted_labels']}")
            print(f"预测标签: {explanation['label_explanations']}")
            print("\n各标签概率:")
            for label, info in explanation['label_explanations'].items():
                status = "✓" if info['predicted'] else "✗"
                print(f"  {status} {label}: {info['probability']:.3f} - {info['description']}")
            print("-" * 80)
        
        # 交互式预测
        print("\n=== 交互式预测 ===")
        print("输入文本进行预测（输入'quit'退出）:")
        
        while True:
            user_input = input("\n请输入文本: ").strip()
            if user_input.lower() == 'quit':
                break
            
            if user_input:
                explanation = classifier.explain_prediction(user_input)
                print(f"\n预测结果: {explanation['predicted_labels']}")
                print("\n详细概率:")
                for label, info in explanation['label_explanations'].items():
                    if info['probability'] > 0.1:  # 只显示概率大于0.1的标签
                        status = "✓" if info['predicted'] else "✗"
                        print(f"  {status} {label}: {info['probability']:.3f}")
    
    except FileNotFoundError:
        print(f"错误: 找不到模型文件 {model_path}")
        print("请先运行 main.py 训练模型")
    except Exception as e:
        print(f"推理过程中出现错误: {e}")

if __name__ == "__main__":
    main()